from langchain.embeddings import OpenAIEmbeddings
from langchain.prompts import PromptTemplate, MaxMarginalRelevanceExampleSelector, FewShotPromptTemplate, \
    SemanticSimilarityExampleSelector
from langchain.vectorstores.faiss import FAISS

examples = [
    {"input": "happy", "output": "sad"},
    {"input": "tall", "output": "short"},
    {"input": "energetic", "output": "lethargic"},
    {"input": "sunny", "output": "gloomy"},
    {"input": "windy", "output": "calm"},
]

example_prompt = PromptTemplate(
    input_variables=["input", "output"],
    template="Input: {input}\nOutput: {output}"
)

example_selector = MaxMarginalRelevanceExampleSelector.from_examples(
    # 这是可供选择的示例列表.
    examples,
    # 这是用于生成嵌入的嵌入类，用于测量语义相似性.
    OpenAIEmbeddings(),
    # 这是 VectorStore 类，用于存储嵌入并执行相似性检索.
    FAISS,
    # 这是要生成的示例数.
    k=1,
)

mmr_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="Give the antonym of every input",
    suffix="Input: {adjective}\nOutput:",
    input_variables=["adjective"],
)

print(mmr_prompt.format(adjective="worried"))

example_selector = SemanticSimilarityExampleSelector.from_examples(
    examples,
    OpenAIEmbeddings(),
    FAISS,
    k=1
)
similar_prompt = FewShotPromptTemplate(
    example_selector=example_selector,
    example_prompt=example_prompt,
    prefix="Give the antonym of every input",
    suffix="Input: {adjective}\nOutput: ",
    input_variables=["adjective"]
)
print(similar_prompt.format(adjective="worried"))